import re
import json

from typing import List, Dict, Any
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser


class InfoNeedManager:
    def __init__(self, cfg: Dict[str, Any]):
        self.retry = cfg.get('retry')

        model = ChatOpenAI(
            model=cfg.get('rag_query_model'),
            base_url=cfg.get('base_url'),
            api_key=cfg.get('api_key'),
            temperature=0.0
        )

        prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                "你是一个信息点抽取助手，**当前正处于 RAG（Retrieval-Augmented Generation，检索增强生成）系统的第一步信息点抽取环节**。\n"
                "整个 RAG 流程分为：\n"
                "1. 信息点抽取（你当前负责，识别需要外部补充的关键信息点）；\n"
                "2. 用这些信息点去知识库（如向量库 FAISS/数据库）检索补充相关内容；\n"
                "3. 检索结果将被下游用于 LLM 生成最终回答。\n"
                "你的任务是：\n"
                "请从对话历史和用户的最新问题中，尽可能多的提取出当前历史记录中没有的、用户可能之前提到过的、大模型不太可能直接知道的、但对当前问题回答非常关键的信息点。\n"
                "比如用户问他自己的生日是什么，你可以返回['用户的生日']；用户说请给我做一个任务画像描写，你可以返回['用户的性格','用户的价值观','用户的年龄'...]。\n"
                "这些信息点包括但不限于：用户的个人属性（如身份证号、住址、生日、项目路径等）、用户最新动态、对话中特有的专有名词、与用户密切相关的事实等。\n"
                "每个信息点请用简洁、明确、具备区分度的表述，并且要标明‘用户的…’或‘用户当前…’等，不能只写泛泛的词语。\n"
                "如果用户指明你去“查一下”或者“去RAG查一下”之类的话，就一定要去查。\n"
                "如果找不到任何此类信息点，请返回一个空的JSON列表：[]。\n"
                "请严格以‘用户’指代对话人，不要使用‘我’、‘你’等其它指代。\n\n"
                "输出格式示例：\n"
                "[\n"
                "  \"用户的AA项目的路径\",\n"
                "  \"用户的银行卡号\",\n"
                "  \"用户的生日\"\n"
                "]\n"
                "以下是最近的对话历史（供你理解背景）：\n"
                "{history}"
            ),
            ("user", "{user_input}")
        ])

        parser = StrOutputParser()

        self.chain = prompt | model | parser

    def extract_info_needs(self, messages: list) -> List[str]:
        user_msg = messages[-1].content if hasattr(messages[-1], "content") else str(messages[-1])

        history_list = []
        for m in messages[:-1]:
            role = "用户" if m.type == "human" or m.__class__.__name__ == "HumanMessage" else "助手"
            history_list.append(f"{role}: {m.content}")
        history_str = "\n".join(history_list)

        for i in range(self.retry):
            raw = self.chain.invoke({"history": history_str, "user_input": user_msg})
            content = raw.strip()
            if content.startswith("```"):
                content = re.sub(r"^```(?:json)?", "", content)
                content = content.rstrip("`").strip()
            content = re.sub(r",\s*([\]}])", r"\1", content)
            try:
                data = json.loads(content)
                if isinstance(data, list):
                    if data:
                        with open("./log.txt", "a", encoding="utf-8") as f:
                            f.write(f"query info: {data}\n")
                    return data
            except json.JSONDecodeError:
                print(f"第{i}次信息点抽取返回非 JSON：\n{raw}")
        print("信息点抽取失败：未返回合法 JSON。")
        return []
